import taichi as ti
import numpy as np
n2=8
n=2**n2
print('n=',n)
ti.init(ti.opengl)

num=ti.field(ti.i32,n)
value=ti.field(ti.f32,n)
result=ti.field(ti.f32,n)
step=ti.field(ti.i32,2)#包含内步长与外步长

step[0]=2#外步长
step[1]=2#内步长

@ti.kernel
def init():
    for i in ti.static(range(n)):
        num[i]=i
@ti.func
def xor(a,b):
    return (a+b)&1
@ti.kernel
def getResult():
    for i in ti.static(range(n)):
        result[i]=value[num[i]]
@ti.kernel
def step1():
    for i in ti.static(range(n//2)):
        halfstep=step[1]//2
        i1=i+(i//(halfstep))*halfstep
        i2=i1+halfstep
        num1,num2=num[i1],num[i2]
        a,b=value[num1],value[num2]
        updown=(i*2//step[0])&1#up,down方向
        if xor(a>b,updown):
            num[i1],num[i2]=num2,num1
def printresult():
    getResult()
    print(num)
    print(value)
    print(result)
init()
value.from_numpy(np.random.rand(n))
print('initial finish')
for i in range(1,n2+1):
    step[0]=2**i
    for j in range(i):
        step[1]=2**(i-j)
        step1()
        print(step[0],step[1])
for i in range(n):
    print(value[num[i]])
printresult()
